{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AtspjpqElmtO"
   },
   "source": [
    "# Inference Speedup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wKxd_dT3lmtQ"
   },
   "source": [
    "## Seminar\n",
    "\n",
    "Some functions, classes and variables required for the homework part are defined in the seminar's code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "axGKC6hXlmtR",
    "outputId": "553dca24-fd13-411a-a322-e7cb48733b10"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from typing import List\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "import triton\n",
    "import triton.language as tl\n",
    "\n",
    "print(\"torch version:\", torch.__version__)\n",
    "print(\"triton version:\", triton.__version__)\n",
    "\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X3NJfn9nlmtS"
   },
   "source": "### Catch activation and weight"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 0,
     "referenced_widgets": [
      "61621d16e2cd40c3ad0c6a849a8b8264",
      "6bed81d4616649bd857418a131538f7d",
      "9ffb3a0c0fea4081bf4f7c5e6218758f",
      "3838c3cdee884f6d902e3b5d0d4e4583",
      "7c3c03fa91874ae5ac941ed171eeb3ea",
      "7498bbbbae7048119584dea4a914da0a",
      "6272889cfea04dab834ee1c62817540b",
      "5f922ecfa23c4705a42336c1e79900df",
      "7f1074b6abde4dddb7e50f67dc9a99b3",
      "87caf0f9f3dc4b508d7733e66fde5407",
      "c0e17b91d9b244daaf63f87f4683503e"
     ]
    },
    "id": "UI2qK8TdlmtT",
    "outputId": "23e582d5-f21a-4fd4-b1d9-13470ba365e3"
   },
   "outputs": [],
   "source": [
    "model_name = \"Qwen/Qwen2.5-3B-Instruct\"\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_name,\n",
    "    torch_dtype=torch.float16,\n",
    "    device_map=\"auto\"\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "streamer = TextStreamer(tokenizer, skip_prompt=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "5GbVOeSXlmtT"
   },
   "outputs": [],
   "source": [
    "# challenging input\n",
    "prompt = \"\"\"\n",
    "# SYSTEM PREAMBLE\n",
    "1) You are an excellent Python software developer with over 10 years of experience. You have a strong understanding of Python related topics, data structures, libraries, frameworks, algorithms, best practices and optimization techniques.\n",
    "2) You are here to help the user (the software developer) by breaking his request in ## TASK into logical steps and writing high-quality and efficient code to implement each step.\n",
    "3) You have to return the entire code.\n",
    "4) Follow \"Answering rules\" without exception.\n",
    "\n",
    "## ANSWERING RULES\n",
    "1) Repeat the question before answering it.\n",
    "2) Always follow \"CHAIN OF THOUGHTS\" to execute the task.\n",
    "\n",
    "## CHAIN OF THOUGHTS\n",
    "1) **OBEY the EXECUTION MODE**\n",
    "2) **TASK ANALYSIS:**\n",
    "   - Understand the user's request thoroughly.\n",
    "   - Identify the key components and requirements of the task.\n",
    "3) **PLANNING: CODDING:**\n",
    "   - Break down the task into logical, sequential steps.\n",
    "   - Outline the strategy for implementing each step.\n",
    "4) **CODING:**\n",
    "   - Explain your thought process before writing any code.\n",
    "   - Write the entire code for each step, ensuring it is clean, optimized, and well-commented.\n",
    "   - Handle edge cases and errors appropriately.\n",
    "5) **VERIFICATION:**\n",
    "   - Review the complete code solution for accuracy and efficiency.\n",
    "   - Ensure the code meets all requirements and is free of errors.\n",
    "\n",
    "## TASK\n",
    "\n",
    "Write a python function that receives the following JSON as input and enters data from it into the Google Sheet.\n",
    "\n",
    "{\n",
    "    'date': '31-05-2024',\n",
    "    'revenue': 90000,\n",
    "    'person' : 'User1',\n",
    "    'expensesList': [30000, 14000, 10000, 2000, 15000],\n",
    "    'expensesDescList': [ 'Ключи', 'Ключи2', 'Счет за такси', 'Клей, пластины', 'Провод 40м'],\n",
    "    'expensesTypeList': ['Закупки', 'Закупки', 'Расходы', 'Ремонт', 'Ремонт']\n",
    "}\n",
    "\n",
    "There is a date in JSON, you can use it to determine the month.\n",
    "The data is entered into a list with the name of the month. If such a list does not exist yet, then you need to create a list with a new month inside the sheet.\n",
    "\n",
    "The list should have the following columns (the first rows are used as headings):\n",
    "A1: Дата расхода,\n",
    "B1: сумма расхода,\n",
    "C1: описание расхода,\n",
    "D1: тип расхода,\n",
    "E1: кто внес данные\n",
    "\n",
    "G1: Дата выручки\n",
    "H1: Сумма выручки\n",
    "I1: Кто внес данные\n",
    "\n",
    "Please separate expenses and profits with a blank column.\n",
    "Please sort expenses by date, including those already listed in Google sheet list.\n",
    "Please sort earnings by date, including those already listed in Google sheet list.\n",
    "\n",
    "It is prohibited to use oauth2client as it is deprecated.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "scDTImShlmtU",
    "outputId": "550cdf9f-6118-46c1-83ea-ab4b3c1059ef"
   },
   "outputs": [],
   "source": [
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"You are Qwen, created by Alibaba Cloud. You are a helpful assistant.\"},\n",
    "    {\"role\": \"user\", \"content\": prompt}\n",
    "]\n",
    "text = tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    tokenize=False,\n",
    "    add_generation_prompt=True\n",
    ")\n",
    "model_inputs = tokenizer([text], return_tensors=\"pt\").to(model.device)\n",
    "generation_output = model.generate(**model_inputs, streamer=streamer, max_new_tokens=1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "GmZ09xWelmtV"
   },
   "outputs": [],
   "source": [
    "class Catcher(nn.Module):\n",
    "    def __init__(self, inps: List, module: nn.Module):\n",
    "        super().__init__()\n",
    "        self.backup = module\n",
    "        self.inps = inps\n",
    "\n",
    "    def forward(self, inp, **kwargs):\n",
    "        self.inps.append(inp.to(\"cpu\"))\n",
    "        raise ValueError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "UaQPJeQMlmtV"
   },
   "outputs": [],
   "source": [
    "layer = model.model.layers[0]\n",
    "inps = []\n",
    "layer.self_attn.q_proj = Catcher(inps, layer.self_attn.q_proj) # wrap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "lDXTyhi4lmtV"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    model(model_inputs.input_ids)\n",
    "except ValueError as e:\n",
    "    layer.self_attn.q_proj = layer.self_attn.q_proj.backup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "N4wTaZR5lmtV",
    "outputId": "9f494935-5a56-49fa-db61-2042d56ba9e7"
   },
   "outputs": [],
   "source": [
    "inps[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "9p1tFHE5lmtV"
   },
   "outputs": [],
   "source": [
    "weight = layer.self_attn.q_proj.weight # unwrap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pHW-5ajhlmtV"
   },
   "source": [
    "### Triton GeMM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### When to Use Triton\n",
    "- Optimization Steps:\n",
    "1. Use torch.compile():\n",
    "    - Start by using torch.compile() to optimize your code.\n",
    "2. Adapt Your Code:\n",
    "    - Rewrite code to be more suitable for torch.compile().\n",
    "        - E.g., eliminate graph breaks to enable CUDA graphs.\n",
    "3. Profile and Identify Bottlenecks:\n",
    "    - Find slow parts of your code using profiling tools.\n",
    "    - Write custom Triton kernels for these parts.\n",
    "4. Consider CUDA:\n",
    "    - If still not fast enough, write custom CUDA kernels.\n",
    "\n",
    "**Note**: For maximum performance from the start, you may choose CUDA directly.\n",
    "\n",
    "#### Rough Edges in Triton\n",
    "- New-ish Project:\n",
    "    - Contains rough edges; code may not behave as expected.\n",
    "    - Expected to become more polished over time.\n",
    "- Recommendation:\n",
    "    - Debugging is important; use “simulator mode” when possible.\n",
    "    - Be aware of limitations on older GPUs or with certain operations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Resources (most useful):\n",
    "- [GPU MODE Lecture 14: Practitioners Guide to Triton](https://christianjmills.com/posts/cuda-mode-notes/lecture-014/#auto-tuning) – basic description about development on Triton, its pros and cons, practical examples with elaborated comments and something about optimization.\n",
    "- [Flash-Decoding for long-context inference](https://pytorch.org/blog/flash-decoding/) – description of SPLIT_K optimization for fasted inference at autoregressive decoding stage due to better GPU utilization.\n",
    "\n",
    "Less useful, but interesting:\n",
    "- [Deep Dive on the Hopper TMA Unit for FP8 GEMMs](https://pytorch.org/blog/hopper-tma-unit/) – importance of TMA unit for Hopper and BlackWell.\n",
    "- [Persistent Matmul](https://triton-lang.org/main/getting-started/tutorials/09-persistent-matmul.html) – special version of GeMM for Hopper and BlackWell with TMA support.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "mH2g_HxalmtV"
   },
   "outputs": [],
   "source": [
    "def is_cuda():\n",
    "    return triton.runtime.driver.active.get_current_target().backend == \"cuda\"\n",
    "\n",
    "assert is_cuda(), \"CUDA only tutorial\"\n",
    "ref_lib = 'cuBLAS'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "I1Il52RRlmtW"
   },
   "outputs": [],
   "source": [
    "def get_cuda_autotune_config():\n",
    "    return [\n",
    "        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),\n",
    "        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 32,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),\n",
    "        triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),\n",
    "\t    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32,  'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "\t    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64,  'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),\n",
    "\t    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),\n",
    "\t    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32,  'GROUP_SIZE_M': 16}, num_stages=4, num_warps=4),\n",
    "\t    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64,  'GROUP_SIZE_M': 16}, num_stages=3, num_warps=8),\n",
    "\t    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),\n",
    "\t    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64,  'GROUP_SIZE_M': 16}, num_stages=4, num_warps=4),\n",
    "\t    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 16}, num_stages=3, num_warps=8),\n",
    "\t    triton.Config({'SPLIT_K': 1, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 32,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32,  'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "\t\ttriton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64,  'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),\n",
    "\t\ttriton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32,  'GROUP_SIZE_M': 16}, num_stages=4, num_warps=4),\n",
    "\t\ttriton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64,  'GROUP_SIZE_M': 16}, num_stages=3, num_warps=8),\n",
    "\t\ttriton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),\n",
    "\t\ttriton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64,  'GROUP_SIZE_M': 16}, num_stages=4, num_warps=4),\n",
    "\t\ttriton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 16}, num_stages=3, num_warps=8),\n",
    "\t\ttriton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),\n",
    "        # custom\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 16,  'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 16,  'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),\n",
    "        triton.Config({'SPLIT_K': 2, 'BLOCK_SIZE_M': 16, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 4, 'BLOCK_SIZE_M': 16, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 4, 'BLOCK_SIZE_M': 16,  'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 4, 'BLOCK_SIZE_M': 16,  'BLOCK_SIZE_N': 64,  'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),\n",
    "        triton.Config({'SPLIT_K': 4, 'BLOCK_SIZE_M': 16, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),\n",
    "        triton.Config({'SPLIT_K': 4, 'BLOCK_SIZE_M': 16, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 16}, num_stages=2, num_warps=4),\n",
    "    ]\n",
    "\n",
    "\n",
    "def get_autotune_config():\n",
    "    if is_cuda():\n",
    "        return get_cuda_autotune_config()\n",
    "    raise NotImplementedError(\"ooops\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "Qx7SUeeZlmtW"
   },
   "outputs": [],
   "source": [
    "@triton.autotune(\n",
    "    configs=get_autotune_config(),\n",
    "    key=['M', 'N', 'K'],\n",
    "    reset_to_zero=['c_ptr']\n",
    ")\n",
    "@triton.jit\n",
    "def matmul_kernel(\n",
    "    # Pointers to matrices\n",
    "    a_ptr, b_ptr, c_ptr,\n",
    "    # Matrix dimensions\n",
    "    M, N, K,\n",
    "    # The stride variables represent how much to increase the ptr by when moving by 1\n",
    "    # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr`\n",
    "    # by to get the element one row down (A has M rows).\n",
    "    stride_am, stride_ak,\n",
    "    stride_bk, stride_bn,\n",
    "    stride_cm, stride_cn,\n",
    "    # Meta-parameters\n",
    "    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,\n",
    "    GROUP_SIZE_M: tl.constexpr, SPLIT_K: tl.constexpr,\n",
    "):\n",
    "    \"\"\"Kernel for computing the matmul C = A x B.\n",
    "    A has shape (M, K), B has shape (K, N) and C has shape (M, N)\n",
    "    \"\"\"\n",
    "    # -----------------------------------------------------------\n",
    "    # Map program ids `pid` to the block of C it should compute.\n",
    "    # This is done in a grouped ordering to promote L2 data reuse.\n",
    "    # See above `L2 Cache Optimizations` section for details.\n",
    "    pid = tl.program_id(axis=0)\n",
    "    pid_sp_k = tl.program_id(axis=1)\n",
    "    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)\n",
    "    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)\n",
    "    num_pid_in_group = GROUP_SIZE_M * num_pid_n\n",
    "    group_id = pid // num_pid_in_group\n",
    "    first_pid_m = group_id * GROUP_SIZE_M\n",
    "    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)\n",
    "    pid_m = first_pid_m + (pid % group_size_m)\n",
    "    pid_n = (pid % num_pid_in_group) // group_size_m\n",
    "\n",
    "    # ----------------------------------------------------------\n",
    "    # Create pointers for the first blocks of A and B.\n",
    "    # We will advance this pointer as we move in the K direction\n",
    "    # and accumulate\n",
    "    # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers\n",
    "    # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers\n",
    "    # See above `Pointer Arithmetics` section for details\n",
    "    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M\n",
    "    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N\n",
    "    offs_k = pid_sp_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)\n",
    "    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)\n",
    "    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)\n",
    "    # -----------------------------------------------------------\n",
    "    # Iterate to compute a block of the C matrix.\n",
    "    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block\n",
    "    # of fp32 values for higher accuracy.\n",
    "    # `accumulator` will be converted back to fp16 after the loop.\n",
    "    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n",
    "    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)):\n",
    "        # Load the next block of A and B, generate a mask by checking the K dimension.\n",
    "        # If it is out of bounds, set it to 0.\n",
    "        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K * SPLIT_K, other=0.0)\n",
    "        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K * SPLIT_K, other=0.0)\n",
    "        # We accumulate along the K dimension.\n",
    "        accumulator += tl.dot(a, b, allow_tf32=False)\n",
    "        # Advance the ptrs to the next K block.\n",
    "        a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak\n",
    "        b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk\n",
    "    # You can fuse arbitrary activation functions here\n",
    "    # while the accumulator is still in FP32!\n",
    "    c = accumulator.to(c_ptr.dtype.element_ty)\n",
    "    # -----------------------------------------------------------\n",
    "    # Write back the block of the output matrix C with masks.\n",
    "    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)\n",
    "    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)\n",
    "    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]\n",
    "    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)\n",
    "    if SPLIT_K == 1:\n",
    "        tl.store(c_ptrs, c, mask=c_mask)\n",
    "    else:\n",
    "        tl.atomic_add(c_ptrs, c, mask=c_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "0OuhzgOylmtW"
   },
   "outputs": [],
   "source": [
    "def matmul(a, b):\n",
    "    # Check constraints.\n",
    "    assert a.shape[1] == b.shape[0], \"Incompatible dimensions\"\n",
    "    M, K = a.shape\n",
    "    K, N = b.shape\n",
    "    # Allocates output.\n",
    "    c = torch.zeros((M, N), device=a.device, dtype=torch.float16)\n",
    "    grid = lambda META: (\n",
    "        triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']),\n",
    "        META['SPLIT_K'],\n",
    "    )\n",
    "    matmul_kernel[grid](\n",
    "        a, b, c,\n",
    "        M, N, K,\n",
    "        a.stride(0), a.stride(1),\n",
    "        b.stride(0), b.stride(1),\n",
    "        c.stride(0), c.stride(1),\n",
    "    )\n",
    "    return c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y7THM_iRlmtW"
   },
   "source": [
    "Sanity check on quality:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Zy7cPBh1lmtW",
    "outputId": "0618602d-24f8-4b83-cf40-175cd154b203"
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "a = inps[0][0].to(torch.float16).cuda()\n",
    "b = weight.to(torch.float16).cuda()\n",
    "\n",
    "triton_output = matmul(a, b)\n",
    "torch_output = torch.matmul(a, b)\n",
    "\n",
    "if torch.allclose(triton_output, torch_output, atol=2e-2):\n",
    "    print(\"✅ Triton and Torch match\")\n",
    "else:\n",
    "    print(\"❌ Triton and Torch differ\")\n",
    "    print(f\"triton_output_with_fp16_inputs={triton_output}\")\n",
    "    print(f\"torch_output_with_fp16_inputs={torch_output}\")\n",
    "    assert False, \"Check quality\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "Wz-10-69lmtW"
   },
   "outputs": [],
   "source": [
    "seqlen = inps[0].shape[1]\n",
    "hidden_size = inps[0].shape[2]\n",
    "seqlen, hidden_size\n",
    "\n",
    "\n",
    "def prepare_a(M: int):\n",
    "    inp = inps[0][0]\n",
    "    if M > seqlen:\n",
    "        n_repeats = M // seqlen + 1\n",
    "        return inp.repeat(n_repeats, 1)[:M]\n",
    "    else:\n",
    "        return inp[:M]\n",
    "\n",
    "\n",
    "def benchmark(M, provider, provider_funcs):\n",
    "    N = weight.shape[0]\n",
    "    K = weight.shape[1]\n",
    "\n",
    "    a = prepare_a(M).to(torch.float16).cuda()\n",
    "    b = weight.T.to(torch.float16)\n",
    "    assert a.shape == (M, K), f\"{a.shape} != {(M, K)}\"\n",
    "    assert b.shape == (K, N), b.shape\n",
    "\n",
    "    quantiles = [0.5, 0.2, 0.8]\n",
    "    ms, min_ms, max_ms = triton.testing.do_bench(lambda: provider_funcs[provider](a, b), quantiles=quantiles)\n",
    "    perf = lambda ms: ms # TFlops = 2 * M * N * K * 1e-12 / (ms * 1e-3)\n",
    "    return perf(ms), perf(max_ms), perf(min_ms)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "H9BhixxhlmtW",
    "outputId": "8e4b0fce-220a-4fe8-c17a-c3377c066954"
   },
   "outputs": [],
   "source": [
    "configs = [\n",
    "    triton.testing.Benchmark(\n",
    "        x_names=[\"M\"],  # Argument names to use as an x-axis for the plot\n",
    "        x_vals=[2 ** i for i in range(7)],  # Different possible values for `x_name`\n",
    "        line_arg=\"provider\",  # Argument name whose value corresponds to a different line in the plot\n",
    "        # Possible values for `line_arg`\n",
    "        line_vals=[ref_lib.lower(), \"triton\"],  # Label name for the lines\n",
    "        line_names=[ref_lib, \"Triton\"],  # Line styles\n",
    "        styles=[(\"green\", \"-\"), (\"blue\", \"-\")],\n",
    "        ylabel=\"ms\",  # Label name for the y-axis\n",
    "        plot_name=\"GeMV-performance-\" + (\"fp16\"),  # Name for the plot, used also as a file name for saving the plot.\n",
    "        args={\"provider_funcs\": {\"triton\": matmul, \"cublas\": torch.matmul}},\n",
    "    ),\n",
    "    triton.testing.Benchmark(\n",
    "        x_names=[\"M\"],  # Argument names to use as an x-axis for the plot\n",
    "        x_vals=[1024 * i for i in range(4, 21)],  # Different possible values for `x_name`\n",
    "        line_arg=\"provider\",  # Argument name whose value corresponds to a different line in the plot\n",
    "        # Possible values for `line_arg`\n",
    "        line_vals=[ref_lib.lower(), \"triton\"],  # Label name for the lines\n",
    "        line_names=[ref_lib, \"Triton\"],  # Line styles\n",
    "        styles=[(\"green\", \"-\"), (\"blue\", \"-\")],\n",
    "        ylabel=\"ms\",  # Label name for the y-axis\n",
    "        plot_name=\"GeMM-performance-\" + (\"fp16\"),  # Name for the plot, used also as a file name for saving the plot.\n",
    "        args={\"provider_funcs\": {\"triton\": matmul, \"cublas\": torch.matmul}},\n",
    "    ),\n",
    "]\n",
    "\n",
    "\n",
    "triton.testing.perf_report(configs)(benchmark).run(show_plots=True, print_data=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2lZEDKGdlmtW"
   },
   "source": [
    "## HW    (10 points total, 4 bonus points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QXEdxgI0lmtW"
   },
   "source": [
    "### Dynamic W8A8 GeMM (6 points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Important disclaimer:** please do the assignment in Google Colab on Nvidia Tesla T4 GPU -- it has all the required dependecies pre-installed and will save you lots of time.\n",
    "\n",
    "For the disobedient: on other GPU types (H100) the code will need to be significantly rewritten to account for architectural features for maximum performance (see \"less useful but interesting\" above in the seminar).\n",
    "\n",
    "> **Important disclaimer 2:** to complete the homework, you need to run the seminar code above, it contains the necessary helper functions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1WLBIKEelmtX"
   },
   "source": [
    "Useful resources:\n",
    "- [Matrix Multiplication Background User's Guide](https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html)\n",
    "- [Deep Dive on CUTLASS Ping-Pong GEMM Kernel](https://pytorch.org/blog/cutlass-ping-pong-gemm-kernel/)\n",
    "- [Accelerating 2D Dynamic Block Quantized Float8 GEMMs in Triton](https://pytorch.org/blog/accelerating-gemms-triton/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://habrastorage.org/webt/xq/r5/8a/xqr58aw0gd6tdm-j45yqad67w1a.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Quantization kernel (2 points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this task you will need to implement fast quantization operation in per-row mode:\n",
    "1. You need to fill in the gaps marked with `YOUR CODE HERE`, please do not delete these comments with the task and marking as it makes checking easier\n",
    "2. Next to each gap there is a comment with `# !!! TASK: ...` that will help understand what exactly needs to be written\n",
    "3. The illustration above shows exactly which dimensions are used to calculate scales for activation and weight matrices\n",
    "4. The formula is `scales = tensor.abs().max(axis=axis) / INT8_max_value`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "1_BCSvnblmtX"
   },
   "outputs": [],
   "source": [
    "@triton.autotune(\n",
    "    configs=[\n",
    "        triton.Config({}, num_stages=2, num_warps=8),\n",
    "        triton.Config({}, num_stages=2, num_warps=4),\n",
    "        triton.Config({}, num_stages=2, num_warps=2),\n",
    "        triton.Config({}, num_stages=2, num_warps=1),\n",
    "     ],\n",
    "    key=['K'],\n",
    ")\n",
    "@triton.jit\n",
    "def quantize_int8_perrow_kernel(\n",
    "    fpa_ptr, a_ptr, as_ptr,\n",
    "    M, K,\n",
    "    stride_fpam, stride_fpak,\n",
    "    stride_am, stride_ak,\n",
    "    stride_asm,\n",
    "    # Meta-parameters\n",
    "    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,\n",
    "):\n",
    "    pid_m = tl.program_id(axis=0)\n",
    "    offs_k = tl.arange(0, BLOCK_SIZE_K)\n",
    "    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M\n",
    "\n",
    "    fpa_ptrs = fpa_ptr + offs_am[:, None] * stride_fpam + offs_k[None, :] * stride_fpak\n",
    "    a_ptrs = a_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak\n",
    "    a_max = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)\n",
    "    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):\n",
    "        fpa = tl.load(fpa_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)\n",
    "\n",
    "        # !!! TASK: calc maximum absolute value of each row of fpa and update a_max\n",
    "        # YOUR CODE HERE\n",
    "        a_max = \n",
    "        # YOUR CODE HERE\n",
    "\n",
    "        fpa_ptrs += BLOCK_SIZE_K * stride_fpak\n",
    "\n",
    "    # !!! TASK: divide a_max by max positive INT8 value\n",
    "    # YOUR CODE HERE\n",
    "    a_scale = \n",
    "    # YOUR CODE HERE\n",
    "\n",
    "    fpa_ptrs = fpa_ptr + offs_am[:, None] * stride_fpam + offs_k[None, :] * stride_fpak\n",
    "    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):\n",
    "        fpa = tl.load(fpa_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)\n",
    "\n",
    "        # !!! TASK: divide fpa by a_scale and convert to INT8\n",
    "        # YOUR CODE HERE\n",
    "        inta = \n",
    "        # YOUR CODE HERE\n",
    "\n",
    "        tl.store(a_ptrs, inta, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K)\n",
    "        fpa_ptrs += BLOCK_SIZE_K * stride_fpak\n",
    "        a_ptrs += BLOCK_SIZE_K * stride_ak\n",
    "    as_offs = pid_m * BLOCK_SIZE_M * stride_asm + tl.arange(0, BLOCK_SIZE_M)\n",
    "    tl.store(as_ptr + as_offs, a_scale)\n",
    "\n",
    "\n",
    "def quantize_int8_perrow(fpa):\n",
    "    a = torch.empty(fpa.shape, device=fpa.device, dtype=torch.int8)\n",
    "    a_scale = torch.empty(fpa.shape[0], device=fpa.device, dtype=fpa.dtype)\n",
    "    M, K = fpa.shape\n",
    "    BLOCK_SIZE_M = 1\n",
    "    BLOCK_SIZE_K = triton.next_power_of_2(K)\n",
    "    grid = (M // BLOCK_SIZE_M,)\n",
    "    quantize_int8_perrow_kernel[grid](\n",
    "        fpa, a, a_scale,\n",
    "        M, K,\n",
    "        fpa.stride(0), fpa.stride(1),\n",
    "        a.stride(0), a.stride(1),\n",
    "        a_scale.stride(0),\n",
    "        BLOCK_SIZE_M, BLOCK_SIZE_K,\n",
    "    )\n",
    "    return a, a_scale\n",
    "\n",
    "\n",
    "def quantize_int8(weight, axis=0, tp_rank=0):\n",
    "    # Weight shape: [H1, H2]\n",
    "    # Scale shape: [H2]\n",
    "    \n",
    "    # !!! TASK: calculate scale by taking maximum over axis and saving dims and divide on maximum positive INT8 value\n",
    "    # YOUR CODE HERE\n",
    "    scale = \n",
    "    # YOUR CODE HERE\n",
    "\n",
    "    # !!! TASK: divide weight by scale and convert to int8\n",
    "    # YOUR CODE HERE\n",
    "    weight = \n",
    "    # YOUR CODE HERE\n",
    "\n",
    "    # col major will accelerate i8xi8 kernel.\n",
    "    if axis == 0:\n",
    "        weight = weight.t().contiguous().t()\n",
    "    scale = scale.squeeze(axis)\n",
    "    return weight.contiguous().cuda(tp_rank), scale.contiguous().cuda(tp_rank)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Per-row W8A8 GeMM (4 points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we need to implement a fast operation for per-row W8A8 GeMM:\n",
    "1. To do this, you need to fill in the gaps marked with `YOUR CODE HERE`. Please do not delete these comments with the task and marking as it makes checking easier\n",
    "2. Next to each gap there is a comment with `# !!! TASK: ...` that will help understand what exactly needs to be written"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Qee3RRE4lmtX"
   },
   "outputs": [],
   "source": [
    "@triton.autotune(\n",
    "    configs=get_autotune_config(),\n",
    "    key=['M', 'N', 'K'],\n",
    "    reset_to_zero=['c_ptr']\n",
    ")\n",
    "@triton.jit\n",
    "def perrow_w8a8_matmul_kernel(\n",
    "    # Pointers to matrices\n",
    "    a_ptr, as_ptr, b_ptr, bs_ptr, c_ptr,\n",
    "    # Matrix dimensions\n",
    "    M, N, K,\n",
    "    # The stride variables represent how much to increase the ptr by when moving by 1\n",
    "    # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr`\n",
    "    # by to get the element one row down (A has M rows).\n",
    "    stride_am, stride_ak,\n",
    "    stride_asm,\n",
    "    stride_bk, stride_bn,\n",
    "    stride_bsn,\n",
    "    stride_cm, stride_cn,\n",
    "    # Meta-parameters\n",
    "    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,\n",
    "    GROUP_SIZE_M: tl.constexpr, SPLIT_K: tl.constexpr,\n",
    "):\n",
    "    \"\"\"Kernel for computing the matmul C = A x B.\n",
    "    A has shape (M, K), B has shape (K, N) and C has shape (M, N)\n",
    "    \"\"\"\n",
    "    # -----------------------------------------------------------\n",
    "    # Map program ids `pid` to the block of C it should compute.\n",
    "    # This is done in a grouped ordering to promote L2 data reuse.\n",
    "    # See above `L2 Cache Optimizations` section for details.\n",
    "    pid = tl.program_id(axis=0)\n",
    "    pid_sp_k = tl.program_id(axis=1)\n",
    "    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)\n",
    "    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)\n",
    "    num_pid_in_group = GROUP_SIZE_M * num_pid_n\n",
    "    group_id = pid // num_pid_in_group\n",
    "    first_pid_m = group_id * GROUP_SIZE_M\n",
    "    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)\n",
    "    pid_m = first_pid_m + (pid % group_size_m)\n",
    "    pid_n = (pid % num_pid_in_group) // group_size_m\n",
    "\n",
    "    # ----------------------------------------------------------\n",
    "    # Create pointers for the first blocks of A and B.\n",
    "    # We will advance this pointer as we move in the K direction\n",
    "    # and accumulate\n",
    "    # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers\n",
    "    # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers\n",
    "    # See above `Pointer Arithmetics` section for details\n",
    "    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M\n",
    "    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N\n",
    "    offs_k = pid_sp_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)\n",
    "    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)\n",
    "    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)\n",
    "    as_ptrs = as_ptr + offs_am * stride_asm\n",
    "    bs_ptrs = bs_ptr + offs_bn * stride_bsn\n",
    "    a_scale = tl.load(as_ptrs, mask=offs_am < M, other=0.0)\n",
    "    b_scale = tl.load(bs_ptrs, mask=offs_bn < N, other=0.0)\n",
    "    # -----------------------------------------------------------\n",
    "    # Iterate to compute a block of the C matrix.\n",
    "    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block\n",
    "    # of fp32 values for higher accuracy.\n",
    "    # `accumulator` will be converted back to fp16 after the loop.\n",
    "    \n",
    "    # !!! TASK: create accumulator of int32 dtype\n",
    "    # YOUR CODE HERE\n",
    "    accumulator = \n",
    "    # YOUR CODE HERE\n",
    "\n",
    "    # !!! TASK: implement the loop over K dimension.\n",
    "    # In the loop first load next blocks of A and B, using an appropriate mask over K dimension.\n",
    "    # Then, multiply blocks and update the accumulator\n",
    "    # Finally, update the pointers to A and B matrix.\n",
    "    # You can rely on implementation of regular matmul from the seminar\n",
    "\n",
    "    # YOUR CODE HERE\n",
    "    for k in range(...):\n",
    "        ...\n",
    "    # YOUR CODE HERE\n",
    "\n",
    "    # You can fuse arbitrary activation functions here\n",
    "    # while the accumulator is still in FP32!\n",
    "    \n",
    "    # !!! TASK: dequantize the accumulator with a_scale and b_scale (outer product) and convert to c_ptr.dtype.element_ty\n",
    "    # YOUR CODE HERE\n",
    "    c = \n",
    "    # YOUR CODE HERE\n",
    "\n",
    "    # -----------------------------------------------------------\n",
    "    # Write back the block of the output matrix C with masks.\n",
    "    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)\n",
    "    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)\n",
    "    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]\n",
    "    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)\n",
    "    if SPLIT_K == 1:\n",
    "        tl.store(c_ptrs, c, mask=c_mask)\n",
    "    else:\n",
    "        tl.atomic_add(c_ptrs, c, mask=c_mask)\n",
    "\n",
    "\n",
    "def matmul_int8(a, a_scale, b, b_scale, out=None):\n",
    "    # Check constraints.\n",
    "    assert a.shape[1] == b.shape[0], \"Incompatible dimensions\"\n",
    "    M, K = a.shape\n",
    "    K, N = b.shape\n",
    "    # Allocates output.\n",
    "    if out == None:\n",
    "        c = torch.zeros((M, N), device=a.device, dtype=torch.float16)\n",
    "    else:\n",
    "        c = out.fill_(0.)\n",
    "    grid = lambda META: (\n",
    "        triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']),\n",
    "        META['SPLIT_K'],\n",
    "    )\n",
    "    perrow_w8a8_matmul_kernel[grid](\n",
    "        a, a_scale, b, b_scale, c,\n",
    "        M, N, K,\n",
    "        a.stride(0), a.stride(1),\n",
    "        a_scale.stride(0),\n",
    "        b.stride(0), b.stride(1),\n",
    "        b_scale.stride(0),\n",
    "        c.stride(0), c.stride(1),\n",
    "    )\n",
    "    return c\n",
    "\n",
    "\n",
    "def matmul_quantize_int8(fpa, b, b_scale, out=None):\n",
    "    # !!! TASK: quantize fpa to int8 and call matmul_int8\n",
    "    # YOUR CODE HERE\n",
    "    a, a_scale = \n",
    "    return \n",
    "    # YOUR CODE HERE\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jtcILz_IlmtX"
   },
   "source": [
    "Sanity check (heavily relaxed for non-smoothed case):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mvuppE_4lmtX",
    "outputId": "56cb7ab9-9bf1-4c6d-c354-a12fd76d2606"
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "cos = torch.nn.CosineSimilarity(0)\n",
    "\n",
    "a = inps[0][0].to(torch.float16).cuda()\n",
    "b = weight.to(torch.float16).cuda()\n",
    "\n",
    "int_a, scale_a = quantize_int8_perrow(a)\n",
    "int_b, scale_b = quantize_int8(b, axis=0)\n",
    "\n",
    "\n",
    "triton_output = matmul_int8(int_a, scale_a, int_b, scale_b)\n",
    "torch_output = torch.matmul(a, b)\n",
    "\n",
    "if (torch_output.float() - triton_output.float()).abs().mean() < 0.03 and torch.quantile((torch_output.float() - triton_output.float()).abs(), 0.95) < 0.07:\n",
    "    print(\"✅ Triton FP8 and Torch match\")\n",
    "else:\n",
    "    print(\"❌ Triton FP8 and Torch differ\")\n",
    "    print(\"Quantization cos: \", cos((int_a * scale_a.unsqueeze(1)).flatten().to(torch.float32), a.flatten().to(torch.float32)).item())\n",
    "\n",
    "    print('=' * 50)\n",
    "    print(f\"triton_output_with_fp8={triton_output}\")\n",
    "    print(f\"torch_output={torch_output}\")\n",
    "\n",
    "    print('=' * 50)\n",
    "    print(\"infs in triton:\", (triton_output).isinf().sum())\n",
    "    print(\"infs in torch:\", (torch_output).isinf().sum())\n",
    "\n",
    "    print('=' * 50)\n",
    "    print(\"Output cos:\", cos(triton_output.flatten().to(torch.float32), torch_output.flatten().to(torch.float32)).item())\n",
    "    print(((triton_output - torch_output).abs() >= 0.5).sum())\n",
    "    print(((triton_output - torch_output).abs() / (torch_output.abs() + 1e-5)))\n",
    "    assert False, \"Triton and Torch differ\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KBItEEG0lmtX"
   },
   "source": [
    "Sanity checks on perf:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "id": "EDGVPDjRlmtX"
   },
   "outputs": [],
   "source": [
    "def test_perf_quantize(M: int, K: int, iters: int = 256, thr: float = 1.3):\n",
    "    torch.manual_seed(0)\n",
    "    print(f\"M: {M} K: {K}\")\n",
    "\n",
    "    a = torch.randn((M, K), device='cuda', dtype=torch.float16)\n",
    "    # warmup\n",
    "    for _ in range(10):\n",
    "        int_a, a_scale = quantize_int8(a, 1)\n",
    "        int_a, a_scale = quantize_int8_perrow(a)\n",
    "    torch.cuda.synchronize()\n",
    "\n",
    "    t1 = time.time()\n",
    "    for _ in range(iters):\n",
    "        int_a, a_scale = quantize_int8_perrow(a)\n",
    "    torch.cuda.synchronize()\n",
    "    t2 = time.time()\n",
    "    for _ in range(iters):\n",
    "        int_a, a_scale = quantize_int8(a, axis=1)\n",
    "    torch.cuda.synchronize()\n",
    "    t3 = time.time()\n",
    "\n",
    "    torch_time = (t3 - t2) / iters\n",
    "    triton_time = (t2 - t1) / iters\n",
    "\n",
    "    print(f\"Torch time cost: {torch_time}\")\n",
    "    print(f\"Triton time cost: {triton_time}\")\n",
    "    assert torch_time / triton_time > thr, f\"Must get at least {thr}x speedup\"\n",
    "    return triton_time, torch_time\n",
    "\n",
    "\n",
    "def test_perf_matmul_int8(M, K, N, iters: int = 512, thr: float = 0.99):\n",
    "    print(\"M: {} K: {} N: {}\".format(M, K, N))\n",
    "    torch.manual_seed(0)\n",
    "    a = torch.randn((M, K), device='cuda', dtype=torch.float16)\n",
    "    b = torch.randn((K, N), device='cuda', dtype=torch.float16).contiguous()\n",
    "    int_b, scale_b = quantize_int8(b, axis=0)\n",
    "    for _ in range(10):\n",
    "        # int_a, a_scale = quantize_int8(a, 1)\n",
    "        int_a, a_scale = quantize_int8_perrow(a)\n",
    "        triton_output = matmul_int8(int_a, a_scale, int_b, scale_b)\n",
    "    torch.cuda.synchronize()\n",
    "    t1 = time.time()\n",
    "    for _ in range(iters):\n",
    "        #int_a, a_scale, _ = quantize_int8(a, 1)\n",
    "        int_a, a_scale = quantize_int8_perrow(a)\n",
    "    torch.cuda.synchronize()\n",
    "    qt2 = time.time()\n",
    "    for _ in range(iters):\n",
    "        triton_output = matmul_int8(int_a, a_scale, int_b, scale_b)\n",
    "    torch.cuda.synchronize()\n",
    "    t2 = time.time()\n",
    "    quant_time = qt2 - t1\n",
    "    triton_time = t2 - qt2\n",
    "    triton_tflops = 2 * M * N * K * 1e-12 / (triton_time / iters)\n",
    "    quant_bandwith = 2 * M * K * 1e-9 / (quant_time / iters)\n",
    "    print(\"Triton time cost: {} (tflops {}) + quant: {} (bandwidth {})\".format(\n",
    "        triton_time, triton_tflops, quant_time, quant_bandwith))\n",
    "    for _ in range(10):\n",
    "        torch_output = torch.matmul(a, b)\n",
    "    torch.cuda.synchronize()\n",
    "    iters = 512\n",
    "    t1 = time.time()\n",
    "    for _ in range(iters):\n",
    "        torch_output = torch.matmul(a, b)\n",
    "    torch.cuda.synchronize()\n",
    "    t2 = time.time()\n",
    "    torch_time = t2 - t1\n",
    "    torch_tflops = 2 * M * N * K * 1e-12 / (torch_time / iters)\n",
    "    print(\"Torch time cost: {} (tflops {})\".format(t2 - t1, torch_tflops))\n",
    "\n",
    "    assert torch_time / triton_time > thr, f\"Must get at least {thr}x speedup\"\n",
    "    return triton_time, torch_time, quant_time\n",
    "\n",
    "\n",
    "def test_perf_model_layer(bs, seq_len, hidden, inter, tp, thr: float = 0.99):\n",
    "    st1 = 0\n",
    "    st2 = 0\n",
    "    st3 = 0\n",
    "    t1, t2, t3 = test_perf_matmul_int8(bs * seq_len, hidden, hidden * 3 // tp, thr=thr)\n",
    "    test_perf_quantize(bs * seq_len, hidden, thr=thr)\n",
    "    st1 += t1\n",
    "    st2 += t2\n",
    "    st3 += t3\n",
    "    t1, t2, t3 = test_perf_matmul_int8(bs * seq_len, hidden // tp, hidden, thr=thr)\n",
    "    test_perf_quantize(bs * seq_len, hidden // tp, thr=thr)\n",
    "    st1 += t1\n",
    "    st2 += t2\n",
    "    st3 += t3\n",
    "    t1, t2, t3 = test_perf_matmul_int8(bs * seq_len, hidden, inter * 2 // tp, thr=thr)\n",
    "    st1 += t1\n",
    "    st2 += t2\n",
    "    st3 += t3\n",
    "    t1, t2, t3 = test_perf_matmul_int8(bs * seq_len, inter // tp, hidden, thr=thr)\n",
    "    test_perf_quantize(bs * seq_len, inter // tp, thr=thr)\n",
    "    st1 += t1\n",
    "    st2 += t2\n",
    "    st3 += t3\n",
    "    print(\"Triton time {} Torch time {} Quant time {}\".format(st1, st2, st3))\n",
    "    assert st2 / st1 > thr, f\"Must get at least {thr}x speedup\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "id": "p8lfPXCJlmtX"
   },
   "outputs": [],
   "source": [
    "bs = 32\n",
    "hidden = model.config.hidden_size\n",
    "inter  = model.config.intermediate_size\n",
    "prefill_len = 512\n",
    "decode_len = 1\n",
    "tp = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TjCbfQOolmtX",
    "outputId": "e12f1a1e-0586-4ff8-f152-e8eefdabf6e5"
   },
   "outputs": [],
   "source": [
    "test_perf_model_layer(bs, prefill_len, hidden, inter, tp, thr=1.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: The last line should show the overall speedup per layer at the prefill stage. For Nvidia Tesla T4 and fixed library versions, our implementation is about 1.6x faster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "SQKwB1pqlmtY",
    "outputId": "5b18ba51-56db-40f6-f202-9536aa7a6340"
   },
   "outputs": [],
   "source": [
    "test_perf_model_layer(bs, decode_len, hidden, inter, tp, thr=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: The last line shows the per-layer comparison at the decode stage. For the same Nvidia Tesla T4 and libraries, the speedup turned out to be much lower (x1.03). To speed up decode, a CUDA implementation is needed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Cmp2zvP7lmtb"
   },
   "source": [
    "Benchmarks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "6JLgqBVblmtb"
   },
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def benchmark_quant_gemm(M, provider, provider_funcs):\n",
    "    N = weight.shape[0]\n",
    "    K = weight.shape[1]\n",
    "\n",
    "    fpa = prepare_a(M).to(torch.float16).cuda().contiguous()\n",
    "    fpb = weight.data.T.to(torch.float16).contiguous()\n",
    "    b, b_scale = quantize_int8(fpb, axis=0)\n",
    "\n",
    "    assert fpa.shape == (M, K), f\"{fpa.shape} != {(M, K)}\"\n",
    "    assert b.shape == (K, N), b.shape\n",
    "\n",
    "    quantiles = [0.5, 0.2, 0.8]\n",
    "    ms, min_ms, max_ms = triton.testing.do_bench(lambda: provider_funcs[provider](fpa, fpb, b, b_scale), quantiles=quantiles)\n",
    "    perf = lambda ms: ms # TFlops = 2 * M * N * K * 1e-12 / (ms * 1e-3)\n",
    "    return perf(ms), perf(max_ms), perf(min_ms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "53fTTwSTlmtc"
   },
   "outputs": [],
   "source": [
    "provider_funcs = {\n",
    "    \"triton int8\": lambda fpa, fpb, b, b_scale: matmul_quantize_int8(fpa, b, b_scale),\n",
    "    \"triton fp16\": lambda fpa, fpb, b, b_scale: matmul(fpa, fpb),\n",
    "    \"cublas fp16\": lambda fpa, fpb, b, b_scale: torch.matmul(fpa, fpb)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "n8BR7sj2lmtc",
    "outputId": "8cff3c7d-0d42-4c31-b796-f46d6fb05b8b"
   },
   "outputs": [],
   "source": [
    "configs = [\n",
    "    triton.testing.Benchmark(\n",
    "        x_names=[\"M\"],  # Argument names to use as an x-axis for the plot\n",
    "        x_vals=[2 ** i for i in range(7)],  # Different possible values for `x_name`\n",
    "        line_arg=\"provider\",  # Argument name whose value corresponds to a different line in the plot\n",
    "        # Possible values for `line_arg`\n",
    "        line_vals=[\"cublas fp16\", \"triton int8\", \"triton fp16\"],  # Label name for the lines\n",
    "        line_names=[\"cuBLAS FP16\", \"Triton INT8\", \"Triton FP16\"],  # Line styles\n",
    "        styles=[(\"green\", \"-\"), (\"blue\", \"-\"), (\"red\", \"-\")],\n",
    "        ylabel=\"ms\",  # Label name for the y-axis\n",
    "        plot_name=\"GeMV-performance-\" + (\"FP16 vs INT8\"),  # Name for the plot, used also as a file name for saving the plot.\n",
    "        args={\"provider_funcs\": provider_funcs},\n",
    "    ),\n",
    "    triton.testing.Benchmark(\n",
    "        x_names=[\"M\"],  # Argument names to use as an x-axis for the plot\n",
    "        x_vals=[1024 * i for i in range(4, 21)],  # Different possible values for `x_name`\n",
    "        line_arg=\"provider\",  # Argument name whose value corresponds to a different line in the plot\n",
    "        # Possible values for `line_arg`\n",
    "        line_vals=[\"cublas fp16\", \"triton int8\", \"triton fp16\"],  # Label name for the lines\n",
    "        line_names=[\"cuBLAS FP16\", \"Triton INT8\", \"Triton FP16\"],  # Line styles\n",
    "        styles=[(\"green\", \"-\"), (\"blue\", \"-\"), (\"red\", \"-\")],\n",
    "        ylabel=\"ms\",  # Label name for the y-axis\n",
    "        plot_name=\"GeMM-performance-\" + (\"FP16 vs INT8\"),  # Name for the plot, used also as a file name for saving the plot.\n",
    "        args={\"provider_funcs\": provider_funcs},\n",
    "    ),\n",
    "]\n",
    "\n",
    "\n",
    "triton.testing.perf_report(configs)(benchmark_quant_gemm).run(show_plots=True, print_data=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** We significantly outperformed Triton FP16 in terms of speed, and even managed to accelerate on prefill with large contexts compared to cuBLAS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "FV3wWlMOlmtc"
   },
   "outputs": [],
   "source": [
    "provider_funcs = {\n",
    "    \"torch\": lambda a, b: quantize_int8(a),\n",
    "    \"triton\": lambda a, b: quantize_int8_perrow(a)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "kPWhpciPlmtc",
    "outputId": "da337105-ec56-4c42-b3d9-e70207fb7513"
   },
   "outputs": [],
   "source": [
    "configs = [\n",
    "    triton.testing.Benchmark(\n",
    "        x_names=[\"M\"],  # Argument names to use as an x-axis for the plot\n",
    "        x_vals=[2 ** i for i in range(7)],  # Different possible values for `x_name`\n",
    "        line_arg=\"provider\",  # Argument name whose value corresponds to a different line in the plot\n",
    "        # Possible values for `line_arg`\n",
    "        line_vals=[\"torch\", \"triton\"],  # Label name for the lines\n",
    "        line_names=[\"Torch\", \"Triton\"],  # Line styles\n",
    "        styles=[(\"green\", \"-\"), (\"blue\", \"-\")],\n",
    "        ylabel=\"ms\",  # Label name for the y-axis\n",
    "        plot_name=\"Quantize-performance-\" + (\"decoding\"),  # Name for the plot, used also as a file name for saving the plot.\n",
    "        args={\"provider_funcs\": provider_funcs},\n",
    "    ),\n",
    "    triton.testing.Benchmark(\n",
    "        x_names=[\"M\"],  # Argument names to use as an x-axis for the plot\n",
    "        x_vals=[256 * i for i in range(4, 21)],  # Different possible values for `x_name`\n",
    "        line_arg=\"provider\",  # Argument name whose value corresponds to a different line in the plot\n",
    "        # Possible values for `line_arg`\n",
    "        line_vals=[\"torch\", \"triton\"],  # Label name for the lines\n",
    "        line_names=[\"Torch\", \"Triton\"],  # Line styles\n",
    "        styles=[(\"green\", \"-\"), (\"blue\", \"-\")],\n",
    "        ylabel=\"ms\",  # Label name for the y-axis\n",
    "        plot_name=\"Quantize-performance-\" + (\"prefill\"),  # Name for the plot, used also as a file name for saving the plot.\n",
    "        args={\"provider_funcs\": provider_funcs},\n",
    "    ),\n",
    "]\n",
    "\n",
    "\n",
    "triton.testing.perf_report(configs)(benchmark).run(show_plots=True, print_data=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** The PyTorch implementation for quantization is extremely inefficient and Triton quantization performs significantly better on both prefill and decode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2R9vd2NOlmtc"
   },
   "source": [
    "### Bonus task: SmoothQuant (4 points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PpxqtWaYlmtc"
   },
   "source": [
    "Your task is to use your implementation of dynamic w8a8 GeMM for real-world LLM inference speedup. Read the paper below and implement the algorithm:\n",
    "1. Smooth your model. You may want to write custom triton kernel for fast activation smoothing in runtime (or leave it in torch).\n",
    "2. After that validate the model quality (e.g. perplexity).\n",
    "3. Create `W8A8Linear` module, which smooths and dynamically quantizes activations and performs quantized GeMM with already smoothed and quantized INT8 weights.\n",
    "4. Replace all linear layers with `W8A8Linear` (quantize your smoothed weight).\n",
    "5. Validate quantized model. Your should experience no drop in quality.\n",
    "6. Measure end-to-end speedup.\n",
    "7. 🤝🏻 Get my personal respect\n",
    "\n",
    "Resources:\n",
    "- https://arxiv.org/abs/2211.10438\n",
    "- [SmoothQuant repo](https://github.com/mit-han-lab/smoothquant)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vanilla SpecDec (4 points)\n",
    "\n",
    "_Special thanks to @hawkeoni for this part of HW._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You might want to reload the notebook. In such case, run this cell.\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "print(\"Using device\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if True:\n",
    "    name_big = \"unsloth/Llama-3.2-3B-Instruct\"\n",
    "    name_small = \"unsloth/Llama-3.2-1B-Instruct\"\n",
    "else:\n",
    "    name_big = \"openai-community/gpt2-medium\"\n",
    "    name_small = \"openai-community/gpt2\"\n",
    "\n",
    "model_3b = AutoModelForCausalLM.from_pretrained(name_big).to(device)\n",
    "model_1b = AutoModelForCausalLM.from_pretrained(name_small).to(device)\n",
    "\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(name_big)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"\n",
    "# SYSTEM PREAMBLE\n",
    "1) You are an excellent Python software developer with over 10 years of experience. You have a strong understanding of Python related topics, data structures, libraries, frameworks, algorithms, best practices and optimization techniques.\n",
    "2) You are here to help the user (the software developer) by breaking his request in ## TASK into logical steps and writing high-quality and efficient code to implement each step.\n",
    "3) You have to return the entire code.\n",
    "4) Follow \"Answering rules\" without exception.\n",
    "\n",
    "## ANSWERING RULES\n",
    "1) Repeat the question before answering it.\n",
    "2) Always follow \"CHAIN OF THOUGHTS\" to execute the task.\n",
    "\n",
    "## CHAIN OF THOUGHTS\n",
    "1) **OBEY the EXECUTION MODE**\n",
    "2) **TASK ANALYSIS:**\n",
    "   - Understand the user's request thoroughly.\n",
    "   - Identify the key components and requirements of the task.\n",
    "3) **PLANNING: CODDING:**\n",
    "   - Break down the task into logical, sequential steps.\n",
    "   - Outline the strategy for implementing each step.\n",
    "4) **CODING:**\n",
    "   - Explain your thought process before writing any code.\n",
    "   - Write the entire code for each step, ensuring it is clean, optimized, and well-commented.\n",
    "   - Handle edge cases and errors appropriately.\n",
    "5) **VERIFICATION:**\n",
    "   - Review the complete code solution for accuracy and efficiency.\n",
    "   - Ensure the code meets all requirements and is free of errors.\n",
    "\n",
    "## TASK\n",
    "\n",
    "Write a python function that receives the following JSON as input and enters data from it into the Google Sheet.\n",
    "\n",
    "{\n",
    "    'date': '31-05-2024',\n",
    "    'revenue': 90000,\n",
    "    'person' : 'User1',\n",
    "    'expensesList': [30000, 14000, 10000, 2000, 15000],\n",
    "    'expensesDescList': [ 'Ключи', 'Ключи2', 'Счет за такси', 'Клей, пластины', 'Провод 40м'],\n",
    "    'expensesTypeList': ['Закупки', 'Закупки', 'Расходы', 'Ремонт', 'Ремонт']\n",
    "}\n",
    "\n",
    "There is a date in JSON, you can use it to determine the month.\n",
    "The data is entered into a list with the name of the month. If such a list does not exist yet, then you need to create a list with a new month inside the sheet.\n",
    "\n",
    "The list should have the following columns (the first rows are used as headings):\n",
    "A1: Дата расхода,\n",
    "B1: сумма расхода,\n",
    "C1: описание расхода,\n",
    "D1: тип расхода,\n",
    "E1: кто внес данные\n",
    "\n",
    "G1: Дата выручки\n",
    "H1: Сумма выручки\n",
    "I1: Кто внес данные\n",
    "\n",
    "Please separate expenses and profits with a blank column.\n",
    "Please sort expenses by date, including those already listed in Google sheet list.\n",
    "Please sort earnings by date, including those already listed in Google sheet list.\n",
    "\n",
    "It is prohibited to use oauth2client as it is deprecated.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Greedy SpecDec (3 points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this task, you need to implement speculative decoding in PyTorch. **Generation must be done greedily.**\n",
    "\n",
    "1. Generate n tokens using the small model\n",
    "2. Check if the big model would choose these same tokens during greedy generation (only one forward pass of the big model should be called, using big_model.generate is not allowed at this stage)\n",
    "3. If all tokens are chosen by the big model, accept them and return to step 1\n",
    "4. If any token is chosen incorrectly, use the correct token from step 2 instead and return to step 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def speculative_generate(big_model, small_model, prefix: str, max_num_tokens: int, n: int) -> tuple[str, int, int]:\n",
    "    \"\"\"Vanilla speculative decoding for greedy generation.\n",
    "\n",
    "    Args:\n",
    "        big_model: original big HF model (verifier).\n",
    "        small_model: small HF model (draft).\n",
    "        prefix (str): prompt.\n",
    "        max_num_tokens (int): max tokens to generate.\n",
    "        n (int): number of tokens to speculate.\n",
    "\n",
    "    Returns: generated text, number of accepted tokens, number of all tokens.\n",
    "    \"\"\"\n",
    "\n",
    "    input_ids = tokenizer(prefix, return_tensors=\"pt\").input_ids.to(device)\n",
    "    start_size = input_ids.size(1)\n",
    "    cnt_accepted = 0\n",
    "    cnt_all = 0\n",
    "    while input_ids.size(1) - start_size < max_num_tokens:\n",
    "        # !!! TASK: Generate a sequence using the small model\n",
    "        small_generation = \n",
    "        num_generated_tokens = \n",
    "\n",
    "        # !!! TASK: Check the sequence using the big model\n",
    "        big_model_logits = \n",
    "        big_model_generations = \n",
    "        \n",
    "        mismatch = False\n",
    "        for i in range(num_generated_tokens):\n",
    "            # !!! TASK: found a mismatch at position i\n",
    "            if True:\n",
    "                mismatch = True\n",
    "                # !!! TASK: if the first token is incorrect, take the first token predicted by the big model\n",
    "                if i == 0:\n",
    "                    input_ids = torch.cat(\n",
    "                        (), \n",
    "                        dim=1\n",
    "                    )\n",
    "                # !!! TASK: otherwise take previous tokens predicted by the small model + the correct token from the big model\n",
    "                else:\n",
    "                    input_ids = torch.cat(\n",
    "                        (),\n",
    "                        dim=1\n",
    "                    )\n",
    "                print(f\"Accepted {i}/{n} tokens\")\n",
    "                cnt_accepted += i\n",
    "                cnt_all += n\n",
    "                break\n",
    "        \n",
    "        if not mismatch:\n",
    "            # !!! TASK: if there was no mismatch, accept the entire sequence + the last token from the big model\n",
    "            input_ids = torch.cat((), dim=1)\n",
    "    # !!! TASK: use tokenizer to decode generated ids\n",
    "    decoded_text = \n",
    "    return decoded_text, cnt_accepted, cnt_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_spec, cnt_accepted, cnt_all = speculative_generate(big_model=model_3b, small_model=model_1b, prefix=prompt, max_num_tokens=128, n=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_inputs = tokenizer([prompt], return_tensors=\"pt\").to(device)\n",
    "res_big = tokenizer.batch_decode(\n",
    "    model_3b.generate(**model_inputs, do_sample=False, max_new_tokens=128)[:, model_inputs.input_ids.shape[1]:],\n",
    "    skip_special_tokens=True\n",
    ")[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert res_big == res_spec and len(res_big) == len(res_spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert cnt_accepted / cnt_all > 0.19"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### HF SpecDec (1 point)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now try using the speculative decoding function from [transformers](https://huggingface.co/docs/transformers/main/en/generation_strategies#speculative-decoding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "# !!! TASK: generate with assistant_model and compare with big model inference time\n",
    "inputs = \n",
    "max_new_tokens = 256\n",
    "\n",
    "# SpecDec from HF\n",
    "start = \n",
    "outputs = \n",
    "print(tokenizer.batch_decode(outputs, skip_special_tokens=True))\n",
    "print(f\"Elapsed time for speculative {}\")\n",
    "\n",
    "# ordinary\n",
    "start = \n",
    "outputs = \n",
    "print(tokenizer.batch_decode(outputs, skip_special_tokens=True))\n",
    "print(f\"Elapsed time for big model inference {}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make a conclusion about the efficiency of SpecDec implementation in HF."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Conslusion:**\n",
    "- TBD"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "2R9vd2NOlmtc"
   ],
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "3838c3cdee884f6d902e3b5d0d4e4583": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_87caf0f9f3dc4b508d7733e66fde5407",
      "placeholder": "​",
      "style": "IPY_MODEL_c0e17b91d9b244daaf63f87f4683503e",
      "value": " 2/2 [00:37&lt;00:00, 17.49s/it]"
     }
    },
    "5f922ecfa23c4705a42336c1e79900df": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "61621d16e2cd40c3ad0c6a849a8b8264": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_6bed81d4616649bd857418a131538f7d",
       "IPY_MODEL_9ffb3a0c0fea4081bf4f7c5e6218758f",
       "IPY_MODEL_3838c3cdee884f6d902e3b5d0d4e4583"
      ],
      "layout": "IPY_MODEL_7c3c03fa91874ae5ac941ed171eeb3ea"
     }
    },
    "6272889cfea04dab834ee1c62817540b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "6bed81d4616649bd857418a131538f7d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_7498bbbbae7048119584dea4a914da0a",
      "placeholder": "​",
      "style": "IPY_MODEL_6272889cfea04dab834ee1c62817540b",
      "value": "Loading checkpoint shards: 100%"
     }
    },
    "7498bbbbae7048119584dea4a914da0a": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "7c3c03fa91874ae5ac941ed171eeb3ea": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "7f1074b6abde4dddb7e50f67dc9a99b3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "87caf0f9f3dc4b508d7733e66fde5407": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "9ffb3a0c0fea4081bf4f7c5e6218758f": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5f922ecfa23c4705a42336c1e79900df",
      "max": 2,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_7f1074b6abde4dddb7e50f67dc9a99b3",
      "value": 2
     }
    },
    "c0e17b91d9b244daaf63f87f4683503e": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
